import subprocess
import sys
import os
from datetime import datetime
import shutil
from pathlib import Path

class PythonEnvironmentManager:
    def __init__(self, python_path):
        self.python_path = self.validate_python_path(python_path)
        self.backup_dir = Path("backups")
        self.backup_dir.mkdir(exist_ok=True)
        self.max_backups = 5
        self.backup_file = None

    @staticmethod
    def validate_python_path(path):
        if not os.path.exists(path):
            raise ValueError(f"Python executable not found at: {path}")
        return path

    def backup_environment(self):
        """备份当前Python环境"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.backup_file = self.backup_dir / f"pip_backup_{timestamp}.txt"
        
        try:
            result = subprocess.run(
                [self.python_path, "-m", "pip", "freeze"],
                capture_output=True,
                text=True,
                check=True
            )
            self.backup_file.write_text(result.stdout)
            print(f"[Backup] Environment saved to {self.backup_file}")
            self.cleanup_old_backups()
            return 0
        except subprocess.CalledProcessError as e:
            print(f"[Error] Backup failed: {e.stderr}")
            self.backup_file = None
            return 1

    def cleanup_old_backups(self):
        """清理旧的备份文件"""
        backups = sorted(self.backup_dir.glob("pip_backup_*.txt"))
        if len(backups) > self.max_backups:
            for backup in backups[:-self.max_backups]:
                backup.unlink()

    def run_command(self, cmd, timeout=300):
        """执行命令，带超时控制"""
        try:
            if isinstance(cmd, str):
                # 如果是 pip install 命令，需要特殊处理路径
                if cmd.strip().startswith('pip install'):
                    parts = cmd.split()
                    # 检查是否是 wheel 文件安装
                    if any(part.endswith('.whl') for part in parts):
                        # 获取当前脚本所在目录
                        script_dir = os.path.dirname(os.path.abspath(__file__))
                        # 修改 wheel 文件路径为绝对路径
                        for i, part in enumerate(parts):
                            if part.endswith('.whl'):
                                parts[i] = os.path.join(script_dir, part)
                        cmd = ' '.join(parts)
                cmd = cmd.split()
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True,
                timeout=timeout,
                cwd=os.path.dirname(os.path.abspath(__file__))  # 设置工作目录为脚本所在目录
            )
            print(result.stdout)
            return result.returncode
        except subprocess.TimeoutExpired:
            print(f"[Error] Command timed out after {timeout} seconds")
            return 1
        except subprocess.CalledProcessError as e:
            print(f"[Error] Command failed: {e.stderr}")
            return 1

    def install_package(self, package_name):
        """安装包，带重试机制"""
        max_retries = 3
        for attempt in range(max_retries):
            try:
                if self.backup_environment() != 0:
                    return 1
                
                cmd = [
                    self.python_path, "-m", "pip", "install",
                    "--prefer-binary", "--force-reinstall",
                    package_name
                ]
                
                return self.run_command(cmd)
            except Exception as e:
                if attempt == max_retries - 1:
                    print(f"[Error] Installation failed after {max_retries} attempts: {e}")
                    return 1
                print(f"[Retry] Attempt {attempt + 1} failed, retrying...")

def main():
    python_path = sys.argv[1] if len(sys.argv) > 1 else "python"
    manager = PythonEnvironmentManager(python_path)
    
    print(f"Python交互式环境(使用: {python_path})")
    print("提示: 直接输入命令执行系统命令")
    print("     使用 > 前缀执行Python代码\n")
    
    while True:
        try:
            user_input = input(">>> ").strip()
            if user_input.lower() in ('exit', 'quit'):
                break
            if not user_input:
                continue
                
            if user_input.startswith('>'):
                try:
                    exec(user_input[1:])
                except Exception as e:
                    print(f"Python执行错误: {e}")
            else:
                manager.run_command(user_input)
                
        except KeyboardInterrupt:
            print("\n退出...")
            sys.exit(0)
        except Exception as e:
            print(f"错误: {e}")

if __name__ == "__main__":
    main()